import os
import json
from datetime import datetime
from pathlib import Path
import subprocess

def load_config():
    with open('config.json', 'r', encoding='utf-8') as f:
        return json.load(f)

def print_message(message):
    print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] {message}")

def get_video_files(video_dir):
    video_extensions = ('.mp4', '.mkv', '.avi', '.mov')
    return [f for f in os.listdir(video_dir) if f.lower().endswith(video_extensions)]

def get_matching_subtitle(video_name, subtitle_dir):
    subtitle_extensions = ('.srt', '.ass', '.ssa')
    video_name_without_ext = os.path.splitext(video_name)[0]
    
    # 处理可能包含目录结构的文件名
    video_name_parts = video_name_without_ext.split(os.sep)
    if len(video_name_parts) > 1:
        # 如果文件名包含路径，只取最后一部分
        video_name_without_ext = video_name_parts[-1]
    
    for ext in subtitle_extensions:
        subtitle_path = os.path.join(subtitle_dir, f"{video_name_without_ext}{ext}")
        if os.path.exists(subtitle_path):
            return subtitle_path
    return None

def merge_video_subtitle(video_path, subtitle_path, output_path):
    # 检查输出文件是否已存在
    if os.path.exists(output_path):
        print_message(f"输出文件已存在，跳过处理: {output_path}")
        return "skipped"
    
    # 检查字幕文件是否存在
    if not os.path.exists(subtitle_path):
        print_message(f"字幕文件不存在: {subtitle_path}")
        return False

    # 处理路径中的特殊字符
    subtitle_path_escaped = subtitle_path.replace('\\', '\\\\').replace(':', '\\:')

    cmd = [
        'ffmpeg',
        '-i', video_path,
        '-vf', f'subtitles=\'{subtitle_path_escaped}\'',
        output_path
    ]
    
    # 输出完整的ffmpeg命令
    print_message(f"执行命令: {' '.join(cmd)}")
    
    try:
        result = subprocess.run(cmd, capture_output=True, text=False)
        if result.returncode != 0:
            error_message = result.stderr.decode('utf-8', errors='ignore')
            print_message(f"Error merging {video_path}: {error_message}")
            return False
        return True
    except Exception as e:
        print_message(f"Exception while merging {video_path}: {str(e)}")
        return False

def main():
    # 加载配置
    config = load_config()
    video_dir = config['video_dir']
    subtitle_dir = config['subtitle_dir']
    
    # 从subtitle_dir路径中提取最后一个目录名作为输出子目录
    output_subdir = Path(subtitle_dir).parts[-1]
    
    # 确保目录存在
    os.makedirs(video_dir, exist_ok=True)
    os.makedirs(subtitle_dir, exist_ok=True)
    output_dir = os.path.join(video_dir, output_subdir)
    os.makedirs(output_dir, exist_ok=True)
    
    print_message("开始处理视频和字幕合并任务")
    
    # 获取所有视频文件
    video_files = get_video_files(video_dir)
    total_videos = len(video_files)
    print_message(f"找到 {total_videos} 个视频文件")
    
    # 处理每个视频文件
    success_count = 0
    for index, video_file in enumerate(video_files, 1):
        print_message(f"处理第 {index}/{total_videos} 个视频: {video_file}")
        
        # 查找匹配的字幕文件
        subtitle_path = get_matching_subtitle(video_file, subtitle_dir)
        if not subtitle_path:
            print_message(f"未找到视频 {video_file} 对应的字幕文件，跳过处理")
            continue
        
        video_path = os.path.join(video_dir, video_file)
        output_path = os.path.join(output_dir, video_file)
        
        print_message(f"合并视频 {video_file} 和字幕 {os.path.basename(subtitle_path)}")
        result = merge_video_subtitle(video_path, subtitle_path, output_path)
        if result == "skipped":
            success_count += 1
        elif result:
            success_count += 1
            print_message(f"成功处理视频: {video_file}")
            
        else:
            print_message(f"处理视频失败: {video_file}")
    
    print_message(f"处理完成。成功: {success_count}/{total_videos}")

if __name__ == '__main__':
    main()